import tushare as ts
import pandas as pd
import numpy as np
from typing import Dict, List, Optional

class MarketSentimentData:
    """市场情绪数据类
    处理市场宽度指标、波动率指标和情绪指标等市场情绪相关数据
    """
    
    def __init__(self, ts_token: str):
        """初始化市场情绪数据类
        
        Args:
            ts_token: Tushare API token
        """
        self.ts_pro = ts.pro_api(ts_token)
    
    def get_market_breadth(self, trade_date: str) -> Optional[Dict]:
        """获取市场宽度指标数据
        
        Args:
            trade_date: 交易日期，格式YYYYMMDD
            
        Returns:
            包含涨跌家数比、新高新低比等数据的字典
        """
        try:
            # 获取当日所有股票的涨跌数据
            df = self.ts_pro.daily(
                trade_date=trade_date
            )
            if df.empty:
                return None
            
            # 计算涨跌家数
            up_count = len(df[df['pct_chg'] > 0])
            down_count = len(df[df['pct_chg'] < 0])
            flat_count = len(df[df['pct_chg'] == 0])
            
            # 计算涨跌比
            up_down_ratio = up_count / down_count if down_count > 0 else float('inf')
            
            # 获取新高新低数据
            high_low = self.ts_pro.hsgt_top10(
                trade_date=trade_date
            )
            if not high_low.empty:
                new_high = len(high_low[high_low['rank'] <= 10])
                new_low = len(high_low[high_low['rank'] >= len(high_low) - 10])
                high_low_ratio = new_high / new_low if new_low > 0 else float('inf')
            else:
                new_high = new_low = high_low_ratio = 0
            
            return {
                'trade_date': trade_date,
                'up_count': up_count,
                'down_count': down_count,
                'flat_count': flat_count,
                'up_down_ratio': up_down_ratio,
                'new_high': new_high,
                'new_low': new_low,
                'high_low_ratio': high_low_ratio
            }
        except Exception as e:
            print(f"获取市场宽度指标数据失败: {e}")
            return None
    
    def get_volatility(self, index_code: str, start_date: str, end_date: str) -> Optional[Dict]:
        """获取波动率指标数据
        
        Args:
            index_code: 指数代码（如：000001.SH为上证指数）
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            
        Returns:
            包含波动率指标数据的字典
        """
        try:
            df = self.ts_pro.index_daily(
                ts_code=index_code,
                start_date=start_date,
                end_date=end_date
            )
            if df.empty:
                return None
            
            # 按日期升序排序
            df = df.sort_values('trade_date')
            
            # 计算日收益率
            df['returns'] = df['close'].pct_change()
            
            # 计算20日波动率
            df['volatility'] = df['returns'].rolling(window=20).std() * np.sqrt(252) * 100
            
            return {
                'trade_date': df['trade_date'].values,
                'close': df['close'].values,
                'volatility': df['volatility'].values
            }
        except Exception as e:
            print(f"获取波动率指标数据失败: {e}")
            return None
    
    def get_arbr(self, stock_code: str, start_date: str, end_date: str, period: int = 26) -> Optional[Dict]:
        """获取人气指标(AR/BR)数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            period: 计算周期，默认26日
            
        Returns:
            包含AR、BR指标数据的字典
        """
        try:
            df = self.ts_pro.daily(
                ts_code=stock_code,
                start_date=start_date,
                end_date=end_date
            )
            if df.empty:
                return None
            
            # 按日期升序排序
            df = df.sort_values('trade_date')
            
            # 计算AR指标：当日开盘价与当日最高价、最低价的关系
            df['ar_high'] = df['high'] - df['open']
            df['ar_low'] = df['open'] - df['low']
            df['ar'] = df['ar_high'].rolling(window=period).sum() / df['ar_low'].rolling(window=period).sum() * 100
            
            # 计算BR指标：最高价、最低价与昨日收盘价的关系
            df['br_high'] = df['high'] - df['close'].shift(1)
            df['br_low'] = df['close'].shift(1) - df['low']
            df['br'] = df['br_high'].rolling(window=period).sum() / df['br_low'].rolling(window=period).sum() * 100
            
            return {
                'trade_date': df['trade_date'].values,
                'ar': df['ar'].values,
                'br': df['br'].values
            }
        except Exception as e:
            print(f"获取人气指标数据失败: {e}")
            return None
    
    def get_cr(self, stock_code: str, start_date: str, end_date: str, period: int = 26) -> Optional[Dict]:
        """获取能量指标(CR)数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            period: 计算周期，默认26日
            
        Returns:
            包含CR指标数据的字典
        """
        try:
            df = self.ts_pro.daily(
                ts_code=stock_code,
                start_date=start_date,
                end_date=end_date
            )
            if df.empty:
                return None
            
            # 按日期升序排序
            df = df.sort_values('trade_date')
            
            # 计算中间价
            df['mid'] = (df['high'] + df['low']) / 2
            
            # 计算CR指标
            df['cr_high'] = df.apply(lambda x: max(0, x['high'] - x['mid'].shift(1)), axis=1)
            df['cr_low'] = df.apply(lambda x: max(0, x['mid'].shift(1) - x['low']), axis=1)
            df['cr'] = df['cr_high'].rolling(window=period).sum() / df['cr_low'].rolling(window=period).sum() * 100
            
            # 计算CR的移动平均
            ma_periods = [5, 10, 20]
            result = {
                'trade_date': df['trade_date'].values,
                'cr': df['cr'].values
            }
            
            for p in ma_periods:
                result[f'cr_ma{p}'] = df['cr'].rolling(window=p).mean().values
            
            return result
        except Exception as e:
            print(f"获取能量指标数据失败: {e}")
            return None